import cv2
import os
import json

def extract_frames(video_path, output_folder, start_time, end_time, interval):
    # 创建输出文件夹
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    frame_count = 0
    current_time = 0
    image_count = 0

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        current_time = frame_count / fps
        if current_time >= start_time and current_time <= end_time:
            # TODO 这里计算频率的地方还有点问题，影响不是太大。
            if (current_time - start_time) % interval < 1 / fps:
                # 保存帧为图像文件
                frame_filename = os.path.join(output_folder, f"frame_{frame_count:04d}.png")
                cv2.imwrite(frame_filename, frame)
                image_count += 1

        frame_count += 1
        if current_time > end_time:
            break

    cap.release()
    print(f"提取了 {image_count} 帧到文件夹 {output_folder}")

def main(config_path):
    # 读取配置文件
    with open(config_path, 'r') as f:
        config = json.load(f)

    video_root = config['video_root']
    output_root = config['output_root']

    for video in config['videos']:
        video_path = os.path.join(video_root, video['path'])
        start_time = video['start_time']
        end_time = video['end_time']
        interval = video['interval']

        video_name = os.path.splitext(os.path.basename(video_path))[0]
        output_folder = os.path.join(output_root, video_name)

        extract_frames(video_path, output_folder, start_time, end_time, interval)

# 示例用法
config_path = 'config.json'
main(config_path)